#ifndef AUTHORTOPIC_ONLINE_CC
#define AUTHORTOPIC_ONLINE_CC
#include<string>
#include<fstream>
#include<sstream>
#include<iostream>
#include<algorithm>
#include<iterator>
#include<vector>
#include<map>
#include<cstdlib>
#include<ctime>
#include<cmath>
#include "def.hpp"
#include "utils.hpp"
#include "topicmodel_lda_abstract.hpp"

class AuthorTopicOnline:public LatentDirichletAllocation{
protected:
  double** var_rho;
  double** var_eta;
  double* var_eta_a;
  double** var_lambda;
  double* var_lambda_k;

  double** cache_digamma_var_eta;
  double* cache_digamma_var_eta_a;
  double** cache_digamma_var_lambda;
  double* cache_digamma_var_lambda_k;

  double** ss_lambda;
  double** ss_eta;

  int var_iter;
  double lowerbound;
  ofstream file_output;

  vector<int> doc_user_stamp; // the discrete user info associated to documents
  map<int,string> user_mapping;
  map<string,int> user_reverse_mapping;
public:
  AuthorTopicOnline();
  virtual void Inference();
  virtual void Prediction();
  virtual double Likelihood() { return 0.0; } ;
  virtual int ReadDOCS();
  virtual int GetDocuments(int); // this indicate how documents are read through into the model
  virtual int SaveModel();
private:
  int Single_E_Step(int);
  int Single_M_Step();
  void Save_Expectation_Phi_VarLambda();
  void Save_Expectation_Theta_VarEta();
};

int AuthorTopicOnline::SaveModel(){
  string phi_file_name = PuzaDEF::Instance()->output_file_name + ".phi";
  file_output.open(phi_file_name.c_str());
  Save_Expectation_Phi_VarLambda();
  file_output.close();
  return 0;
}


AuthorTopicOnline::AuthorTopicOnline(){
  allocate_memory<double>(var_eta,PuzaDEF::Instance()->AUTHOR_NUM,PuzaDEF::Instance()->TOPIC_NUM);
  allocate_memory<double>(ss_eta,PuzaDEF::Instance()->AUTHOR_NUM,PuzaDEF::Instance()->TOPIC_NUM);
  allocate_memory<double>(cache_digamma_var_eta,PuzaDEF::Instance()->AUTHOR_NUM,PuzaDEF::Instance()->TOPIC_NUM);
  allocate_memory<double>(var_lambda, PuzaDEF::Instance()->TOPIC_NUM, PuzaDEF::Instance()->TERM_NUM);
  allocate_memory<double>(ss_lambda, PuzaDEF::Instance()->TOPIC_NUM, PuzaDEF::Instance()->TERM_NUM);
  allocate_memory<double>(cache_digamma_var_lambda, PuzaDEF::Instance()->TOPIC_NUM, PuzaDEF::Instance()->TERM_NUM);

  
  cache_digamma_var_eta_a = new double[PuzaDEF::Instance()->AUTHOR_NUM];
  cache_digamma_var_lambda_k = new double[PuzaDEF::Instance()->TOPIC_NUM];
  var_lambda_k = new double[PuzaDEF::Instance()->TOPIC_NUM];
  var_eta_a = new double[PuzaDEF::Instance()->AUTHOR_NUM];

  // initialize
  for(int k=0; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
    var_lambda_k[k] = 0.0;
    for(int v=0; v < PuzaDEF::Instance()->TERM_NUM; v++){
      var_lambda[k][v] =  random_gamma_distribution(100,0.01);
      var_lambda_k[k] = var_lambda_k[k] + var_lambda[k][v];

      cache_digamma_var_lambda[k][v] = digamma(var_lambda[k][v]);
    }
    cache_digamma_var_lambda_k[k] = digamma(var_lambda_k[k]);
  }

  for(int a=0; a < PuzaDEF::Instance()->AUTHOR_NUM; a++){
    var_eta_a[a] = 0.0;
    for(int k=0; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
      var_eta[a][k] = random_gamma_distribution(100,0.01);
      var_eta_a[a] = var_eta_a[a] + var_eta[a][k];

      cache_digamma_var_eta[a][k] = digamma(var_eta[a][k]);
    }
    cache_digamma_var_eta_a[a] = digamma(var_eta_a[a]);
  }
}


int AuthorTopicOnline::Single_E_Step(int d){
  TermList current_list = TM_DOCS[d];
  int term_count = current_list.size();

  // E-step for one single document                                                                                                                                         
  double var_rho_sum = 0;
  int current_author = doc_user_stamp[d];
  for(int i=0; i < term_count; i++){
    int term_id = current_list[i].term_id;
    var_rho_sum = 0;
    // calculate un-normalized rho
    for(int k=0; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
      var_rho[i][k] = cache_digamma_var_lambda[k][term_id] - cache_digamma_var_lambda_k[k] + cache_digamma_var_eta[current_author][k] - cache_digamma_var_eta_a[current_author]; // in log space
      if(k == 0){
	var_rho_sum = var_rho[i][k];
      }
      else{
	var_rho_sum = log_sum(var_rho_sum,var_rho[i][k]);
      }
    }
      // normalize rho                                                                                                                                                         
      for(int k=0; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
        var_rho[i][k] = exp(var_rho[i][k] - var_rho_sum);
      }
  }

  // gather statistics and update lower bound                                                          
  for(int k=0; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
    for(int i=0; i < term_count; i++){
      int term_id = current_list[i].term_id;
      int count = current_list[i].term_stat;
      ss_lambda[k][term_id] = ss_lambda[k][term_id] + count * var_rho[i][k];
      ss_eta[current_author][k] = ss_eta[current_author][k] + count * var_rho[i][k];

      lowerbound = lowerbound + count * var_rho[i][k] * (cache_digamma_var_lambda[k][term_id] - cache_digamma_var_lambda_k[k]);
      lowerbound = lowerbound + count * var_rho[i][k] * (cache_digamma_var_eta[current_author][k] - cache_digamma_var_eta_a[current_author]);
      lowerbound = lowerbound - count * var_rho[i][k] * log(var_rho[i][k]);
    }
  }
  return 0;
}


void AuthorTopicOnline::Inference(){
  allocate_memory<double>(var_rho, MAX_INT, PuzaDEF::Instance()->TOPIC_NUM);
  PuzaLogger::Instance()->PutString("Start Batch Variational Inference for Author Topic Model");
  // variational EM                                                                                                             
  var_iter = 0;
  double old_lowerbound = 0;
  while(1){
    lowerbound = 0.0;
    // E-step
    for(int d=0; d < PuzaDEF::Instance()->DOC_NUM; d++){
      Single_E_Step(d);
    }
    // M-step
    Single_M_Step();
    if(var_iter % 20 == 0){
      string phi_file_name = PuzaDEF::Instance()->output_file_name + ".phi";
      file_output.open(phi_file_name.c_str());
      Save_Expectation_Phi_VarLambda();
      file_output.close();
    }
    PuzaLogger::Instance()->PutString("iter :" + to_string<int>(var_iter,std::dec) +  " lowerbound:" + to_string<double>(lowerbound,std::dec));
    if(var_iter > 100){
      if((lowerbound - old_lowerbound)<=0.00001)
        break;
    }
    if(var_iter >= 500)
      break;

    old_lowerbound = lowerbound;
    var_iter ++;
  }
  PuzaLogger::Instance()->PutString("Batch Learning Finished.");
}


int AuthorTopicOnline::Single_M_Step(){
  // update lambda
  for(int k=0; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
    // update lower bound before updating new values                                                                                                                       
    for(int v=0; v < PuzaDEF::Instance()->TERM_NUM; v++){
      lowerbound = lowerbound + (beta[v]-1) * (cache_digamma_var_lambda[k][v] - cache_digamma_var_lambda_k[k]);
      lowerbound = lowerbound - log_gamma(beta[v]);

      lowerbound = lowerbound - (var_lambda[k][v]-1) * (cache_digamma_var_lambda[k][v] - cache_digamma_var_lambda_k[k]);
      lowerbound = lowerbound + log_gamma(var_lambda[k][v]);
    }
    lowerbound = lowerbound + log_gamma(BetaSum);
    lowerbound = lowerbound - log_gamma(var_lambda_k[k]);

    var_lambda_k[k] = 0.0;
    for(int v=0; v < PuzaDEF::Instance()->TERM_NUM; v++){
      var_lambda[k][v] = beta[v] + ss_lambda[k][v];
      var_lambda_k[k] = var_lambda_k[k] + var_lambda[k][v];

      ss_lambda[k][v] = 0.0;
      cache_digamma_var_lambda[k][v] = digamma(var_lambda[k][v]);
    }
    cache_digamma_var_lambda_k[k] = digamma(var_lambda_k[k]);
  }

  // update eta
  for(int a=0; a < PuzaDEF::Instance()->AUTHOR_NUM; a++){
    // update lower bound                                                                                                               
    for(int k=0; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
      lowerbound = lowerbound + (alpha[k]-1) * (cache_digamma_var_eta[a][k] - cache_digamma_var_eta_a[a]);
      lowerbound = lowerbound - log_gamma(alpha[k]);
      lowerbound = lowerbound - (var_eta[a][k]-1) * (cache_digamma_var_eta[a][k] - cache_digamma_var_eta_a[a]);
      lowerbound = lowerbound + log_gamma(var_eta[a][k]);
    }
    lowerbound = lowerbound + log_gamma(AlphaSum);
    lowerbound = lowerbound - log_gamma(var_eta_a[a]);

    var_eta_a[a] = 0.0;
    for(int k=0; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
      var_eta_a[k] = alpha[k] + ss_eta[a][k];
      var_eta_a[a] = var_eta_a[a] + var_eta[a][k];

      ss_eta[a][k] = 0.0;
      cache_digamma_var_eta[a][k] = digamma(var_eta[a][k]);
    }
    cache_digamma_var_eta_a[a] = digamma(var_eta_a[a]);
  }
  return 0;
}


void AuthorTopicOnline::Save_Expectation_Phi_VarLambda(){
  for(int t = 0; t < PuzaDEF::Instance()->TERM_NUM; t++){
    file_output << t << "\t";
    for(int k=0; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
      file_output << k << ":" << (var_lambda[k][t] / var_lambda_k[k]) << " ";
    }
    file_output << endl;
  }
}

void AuthorTopicOnline::Save_Expectation_Theta_VarEta(){
  for(int r=0; r < PuzaDEF::Instance()->AUTHOR_NUM; r++){
    file_output << user_mapping[r] << "\t";
    for(int k=0; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
      file_output << k << ":" << (var_eta[r][k] / var_eta_a[r]) << " ";
    }
    file_output << endl;
  }
}


void AuthorTopicOnline::Prediction(){
}

int AuthorTopicOnline::ReadDOCS(){
  MAX_INT = 0;
  string temp;
  if(PuzaDEF::Instance()->input_file_name.empty()){
    cout << "Please give input file name." << endl;
    exit(0);
  }
  PuzaDEF::Instance()->InputFile.open(PuzaDEF::Instance()->input_file_name.c_str());
  if(!PuzaDEF::Instance()->InputFile){
    cout << "Can't open the file." << endl;
    exit(0);
  }
  PuzaDEF::Instance()->data_ids.clear();
  TM_DOCS.clear();
  string original_id;
  string temp_user_stamp;
  vector<string> buffer;
  while(!PuzaDEF::Instance()->InputFile.eof()){
    getline(PuzaDEF::Instance()->InputFile,temp,'\n');
    std::istringstream iss(temp);
    // get the doc id
    getline(iss,temp,'\t');
    original_id = temp;
    // get user
    getline(iss,temp,'\t');
    temp_user_stamp = temp;
    // get the doc body                                                                                                                                                 
    getline(iss,temp,'\n');
    std::istringstream tempiss(temp);
    vector<string> tokens;
    copy(istream_iterator<string>(tempiss), istream_iterator<string>(), back_inserter<vector<string> >(tokens));
    if(tokens.size()<1){
      continue;
    }

    int now_id = PuzaDEF::Instance()->data_ids.size();
    int user_id = -1;
    PuzaDEF::Instance()->data_ids[now_id] = original_id;

    map<string,int>::iterator user_iter = user_reverse_mapping.find(temp_user_stamp);
    if(user_iter != user_reverse_mapping.end()){
      // id found
      user_id = (*user_iter).second;
      doc_user_stamp.push_back(user_id);
    }
    else{
      user_id = user_reverse_mapping.size();
      user_reverse_mapping[temp_user_stamp] = user_id;
      user_mapping[user_id] = temp_user_stamp;
      doc_user_stamp.push_back(user_id);
    }
    TermList assignList;
    for (vector<string>::iterator iter = tokens.begin(); iter < tokens.end(); iter++){
      int word_id;
      int word_count;
      std::istringstream valueiss((*iter));
      getline(valueiss,temp,':');
      from_string<int>(word_id,temp,std::dec);
      getline(valueiss,temp);  // get count
      from_string<int>(word_count,temp,std::dec);
      TermItem newAssign;
      newAssign.term_id = word_id;
      newAssign.term_stat = word_count;
      assignList.push_back(newAssign);
    }
    if((int)assignList.size() > MAX_INT)
      MAX_INT = assignList.size();
    TM_DOCS.push_back(assignList);
  }
  PuzaDEF::Instance()->InputFile.close();
  return 0;
}

int AuthorTopicOnline::GetDocuments(int choice){
  if(choice == 0){
    return ReadDOCS();
  }
  else{
    PuzaDEF::Instance()->BatchSize = -1;
    return ReadDOCS();
  }
}


#endif
